package com.jorado.pulsar.client.trace;

import brave.Span;
import brave.internal.HexCodec;
import brave.propagation.TraceContext;
import com.jorado.json.JsonUtils;
import com.jorado.trace.util.RequestID;
import org.apache.pulsar.client.api.Consumer;
import org.apache.pulsar.client.api.ConsumerInterceptor;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.impl.TopicMessageImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;

public class TracingConsumerInterceptor<T> implements ConsumerInterceptor<T> {

    private static final Logger LOG = LoggerFactory.getLogger(TracingConsumerInterceptor.class);
    private String clientAppName = "pulsar";
    private static ThreadLocal<Long> beforeTime = new ThreadLocal<>();

    private TraceContextContainer traceContextContainerForBeforeConsume;

    private TraceContextContainer traceContextContainerForAcknowledge;

    public TracingConsumerInterceptor(String clientAppName) {
        traceContextContainerForBeforeConsume = new TraceContextContainer("consume", clientAppName);
        traceContextContainerForAcknowledge = new TraceContextContainer("acknowledge", clientAppName);
        this.clientAppName = clientAppName;
    }


    @Override
    public void close() {

    }

    @Override
    public Message<T> beforeConsume(Consumer<T> consumer, Message<T> message) {
        long spanId = 0L;
        try {
            spanId = Long.parseLong(message.getProperty("spanId"));
        } catch (Exception ignore) {
        }
        boolean sample = true;
        try {
            sample = Boolean.parseBoolean(message.getProperty("sample"));
        } catch (Exception ignore) {
        }
        TraceContext traceContext = traceContextContainerForBeforeConsume.initZipkinContext(message.getProperty("requestId"),
                spanId, sample, true);
        writeConsumeLog(consumer, message, traceContext);
        beforeTime.set(System.currentTimeMillis());
        return message;
    }

    @Override
    public void onAcknowledge(Consumer<T> consumer, MessageId messageId, Throwable throwable) {
        Span currentSpan = traceContextContainerForAcknowledge.getTracer().currentSpan();
        if (currentSpan != null) {
            TraceContext current = traceContextContainerForAcknowledge.getTracer().currentSpan().context();
            TraceContext traceContext = traceContextContainerForAcknowledge.initZipkinContext(RequestID.getRequestID(),
                    current.parentId(), current.sampled(), true);
            writeAckLog(consumer, messageId, throwable, traceContext);
        }
    }

    @Override
    public void onAcknowledgeCumulative(Consumer<T> consumer, MessageId messageId, Throwable throwable) {
    }

    @Override
    public void onNegativeAcksSend(Consumer<T> consumer, Set<MessageId> messageIds) {

    }

    @Override
    public void onAckTimeoutSend(Consumer<T> consumer, Set<MessageId> messageIds) {

    }

    protected void writeConsumeLog(Consumer consumer, Message<T> message, TraceContext traceContext) {
        Map invLog = new HashMap<>();
        invLog.put("type", "ss");
        invLog.put("service", "ConsumerConsumeClient");
        invLog.put("method", "beforeConsume");
        invLog.put("ExecuteTime", 0);
        String topicName;
        if (message instanceof TopicMessageImpl) {
            topicName = ((TopicMessageImpl<byte[]>) message).getTopicName();
        } else {
            topicName = consumer.getTopic();
        }
        invLog.put("topicName", topicName);
        invLog.put("serverType", "httpService");
        invLog.put("Params", "");
        invLog.put("project", "pulsar");
        invLog.put("messageId", message.getMessageId().toString());
        invLog.put("spanId", HexCodec.toLowerHex(traceContext.spanId()));
        invLog.put("parentId", traceContext.parentId() == null ? null : HexCodec.toLowerHex(traceContext.parentId()));
        LOG.info(JsonUtils.toJson(invLog));
    }

    protected void writeAckLog(Consumer<T> consumer, MessageId messageId, Throwable throwable, TraceContext traceContext) {
        Map invLog = new HashMap<>();
        invLog.put("type", "ss");
        invLog.put("service", "ConsumerConsumeClient");
        invLog.put("method", "onAcknowledge");
        invLog.put("project", clientAppName);
        invLog.put("pulsarTags", "pulsar");
        invLog.put("Params", "");
        invLog.put("serverType", "httpService");
        invLog.put("ExecuteTime", beforeTime.get() == null ? null : System.currentTimeMillis() - beforeTime.get());
        invLog.put("topicName", consumer.getTopic());
        invLog.put("messageId", messageId.toString());
        invLog.put("spanId", HexCodec.toLowerHex(traceContext.spanId()));
        if (traceContext.parentId() != null) {
            invLog.put("parentId", HexCodec.toLowerHex(traceContext.parentId()));
        }
        if (throwable != null) {
            invLog.put("errorCode", throwable);
        } else {
            invLog.put("errorCode", "success");
        }
        LOG.info(JsonUtils.toJson(invLog));
    }

    public TraceContextContainer getTraceContextContainerForBeforeConsume() {
        return traceContextContainerForBeforeConsume;
    }

    public TraceContextContainer getTraceContextContainerForAcknowledge() {
        return traceContextContainerForAcknowledge;
    }
}
